package com.blog.blog.util.aop;

import com.blog.blog.common.ApiBaseResponse;
import com.blog.blog.util.RedisRateLimiter;
import com.blog.blog.util.annotations.RateLimit;
import jakarta.servlet.http.HttpServletRequest;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.Objects;

@Aspect
@Component
public class RateLimitAspect {

    private final RedisRateLimiter rateLimiter;

    public RateLimitAspect(RedisRateLimiter rateLimiter) {
        this.rateLimiter = rateLimiter;
    }

    @Around("@annotation(com.blog.blog.util.annotations.RateLimit)")
    public Object around(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        RateLimit anno = method.getAnnotation(RateLimit.class);

        String key = anno.key();
        if (key == null || key.isBlank()) {
            key = method.getDeclaringClass().getSimpleName() + ":" + method.getName();
        }

        if (anno.byIp()) {
            ServletRequestAttributes attrs = (ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes());
            HttpServletRequest req = attrs.getRequest();
            String ip = getClientIp(req);
            key = key + ":ip:" + ip;
        }

        boolean pass = rateLimiter.tryAcquire(key, anno.rate(), anno.interval(), anno.unit());
        if (!pass) {
            return ResponseEntity.status(429).body(ApiBaseResponse.error(429, "请求过于频繁，请稍后重试"));
        }
        return pjp.proceed();
    }

    private String getClientIp(HttpServletRequest request) {
        String ip = request.getHeader("X-Forwarded-For");
        if (ip != null && !ip.isBlank()) {
            int idx = ip.indexOf(',');
            if (idx > 0) ip = ip.substring(0, idx);
            return ip.trim();
        }
        ip = request.getHeader("X-Real-IP");
        if (ip != null && !ip.isBlank()) return ip;
        return request.getRemoteAddr();
    }
}


